"""
python fujian2_RandomForest_predict.py
"""

import os
import re
import json
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
import matplotlib.pyplot as plt
from datetime import datetime, timedelta

def read_json_data(file_path):
    with open(file_path, 'r') as f:
        data = json.load(f)
    df = pd.DataFrame(data)
    df['date'] = pd.to_datetime(df['date'])
    df['sales'] = df['sales'].astype(int)
    return df

def prepare_data(df):
    # 对日期进行排序
    df = df.sort_values('date')
    # 提取特征和目标
    X = df.index.values.reshape(-1, 1)
    y = df['sales'].values
    return X, y

def train_random_forest(X_train, y_train):
    rf = RandomForestRegressor(n_estimators=100, random_state=42)
    rf.fit(X_train, y_train)
    return rf


def predict_future_sales(model, last_date, num_days):
    future_dates = [last_date + timedelta(days=i) for i in range(1, num_days + 1)]
    future_dates_df = pd.DataFrame({'date': future_dates})
    
    # 确保索引的计算正确
    future_dates_df['index'] = future_dates_df.index.values.reshape(-1, 1)  # 直接使用索引
    X_future = future_dates_df['index'].values.reshape(-1, 1)
    
    predictions = model.predict(X_future)  # 确保这个返回的长度等于 num_days
    return future_dates_df, predictions


def save_predictions_to_json(category_id, predictions, output_folder):
    future_dates, predicted_sales = predictions
    result = [{'date': date.strftime('%Y/%m/%d'), 'predicted_sales': int(sale)} for date, sale in zip(future_dates['date'], predicted_sales)]
    output_json_path = os.path.join(output_folder, f'json_category{category_id}.json')
    os.makedirs(os.path.dirname(output_json_path), exist_ok=True)
    with open(output_json_path, 'w') as f:
        json.dump(result, f)

def plot_data_and_predictions(df, predictions, category_id, output_folder):
    original_dates = df['date']
    original_sales = df['sales']
    future_dates, predicted_sales = predictions
    plt.figure(figsize=(10, 6))
    plt.scatter(original_dates, original_sales, label='Original Data', c='blue')
    # plt.scatter(future_dates, predicted_sales, label='Predicted Data', c='red')
    plt.scatter(future_dates['date'], predicted_sales, label='Predicted Data', c='red')
    
    plt.legend()
    plt.xlabel('Date')
    plt.ylabel('Sales')
    plt.title(f'Random Forest Prediction on Category {category_id}')
    output_image_path = os.path.join(output_folder, 'gragh', f'gragh_category{category_id}.png')
    os.makedirs(os.path.dirname(output_image_path), exist_ok=True)
    plt.savefig(output_image_path)
    plt.close()  # 关闭当前图形


if __name__ == "__main__":
    # 输入目录和输出目录
    input_folder = r'..\fujian\fujian2\groupByCategory'
    # output_folder = r'./test'
    output_folder = r'..\fujian\fujian2\RandomForest'
    
    json_files = [os.path.join(input_folder, file) for file in os.listdir(input_folder) if file.endswith('.json')]
    category_id = 1
    for json_file in json_files:
        
        filename = os.path.basename(json_file)
        # print("文件名:", filename)

        # 使用正则表达式提取 category 后的数字
        category_id = -1
        match = re.search(r'category_(?:category)?(\d+)', filename)
        if match:
            category_id = int(match.group(1))
            # print("提取的 category ID:", category_id)
        else:
            print("未找到匹配的 category ID")
        
        # 读取数据
        df = read_json_data(json_file)
        if len(df) < 5:  # 确保样本数量足够
            print(f"数据样本不足: {len(df)}, category的id是: ", category_id)
            continue
        # print(df)
        # 准备数据
        X, y = prepare_data(df)
        # print(X, y)
        # 划分训练集和测试集（这里简单地按 80% 和 20% 划分）
        split_index = int(0.8 * len(df))
        X_train, X_test = X[:split_index], X[split_index:]
        y_train, y_test = y[:split_index], y[split_index:]
        # 训练随机森林模型
        # print(X_train, y_train)
        model = train_random_forest(X_train, y_train)
        # print(model)
        
        # 预测未来销售
        last_date = df['date'].max()
        future_predictions = predict_future_sales(model, last_date, 92)
        # 保存预测到 JSON 文件
        save_predictions_to_json(category_id, future_predictions, output_folder)
        # 绘制原始数据和预测数据的散点图
        plot_data_and_predictions(df, future_predictions, category_id, output_folder)
        category_id += 1
